'''
Author: WangNing
Date: 2025-04-26 10:28:56
LastEditors: WangNing
LastEditTime: 2025-04-27 10:29:55
FilePath: \linetest\backend\api\ai_routes.py
Description: AI服务API路由
Copyright (c) 2025 by VGE, All Rights Reserved. 
'''

from flask import Blueprint, request, Response, g, stream_with_context
import json
from services.ai_service import ai_stream_response

# 创建蓝图
ai_bp = Blueprint('ai', __name__, url_prefix='/api/ai')

@ai_bp.route('/stream', methods=['POST'])
def stream_ai_response():
    """处理AI流式响应请求，返回流式JSON响应"""
    try:
        # 获取请求数据
        data = request.get_json()
        
        if not data or 'prompt' not in data:
            return Response(json.dumps({
                'type': 'error',
                'message': '未提供有效的prompt'
            }), mimetype='application/json')
        
        prompt = data.get('prompt')
        username = data.get('username', 'admin')
        
        # 定义流式响应生成器
        def generate():
            try:
                for chunk in ai_stream_response(prompt, g.db, username):
                    yield f"{chunk}\n"  # 每个JSON响应后加换行符分隔
            except Exception as e:
                yield json.dumps({
                    'type': 'error',
                    'message': f'生成响应时出错: {str(e)}'
                })
                
        # 使用stream_with_context包装生成器，确保请求上下文在整个流式响应过程中可用
        return Response(
            stream_with_context(generate()),
            mimetype='application/json'
        )
    
    except Exception as e:
        return Response(json.dumps({
            'type': 'error',
            'message': f'处理请求时出错: {str(e)}'
        }), mimetype='application/json') 